#include "comm/znx_string.h"
#include "comm/znx_inet.h"
#include "comm/znx_random.h"
#include "event/znx_conn.h"
#include "event/znx_tcp_proc.h"
#include "event/znx_udp_proc.h"
#include "event/znx_worker_proc.h"
#include "event/znx_session_pool.h"

// Five scenarios were tested
// 1. The idle connection exceeds the upper limit of the connection pool.
// 2. The idle connection exceeds the upper limit of a single host in the connection pool
// 3. The connection put into the connection pool was closed due to timeout
// 4. The server waits for a timeout and actively shuts down
// 5. proc_exit, close all idle connections.


static void
znx_test_server_read(znx_event_t *event)
{
    znx_conn_t *server_conn = event->data;
    char data[512];
    size_t data_len = 0;
    znx_result_t ret;

    while (1) {
        ret = server_conn->operator->read(server_conn, data, 500);
        if (ret > 0) {
            data[ret] = '\0';
            data_len = (size_t)ret;
            ZNX_DEF_LOG_DEBUG(NULL, "test_server read data: %s", data);
            break;
        }

        if (ret == ZNX_AGAIN) {
            ZNX_DEF_LOG_DEBUG(NULL, "test_server read EAGAIN");
            znx_event_loop_enable_read(server_conn);
            return;
        }

        ZNX_DEF_LOG_DEBUG(NULL, "test_server: client close");
        znx_event_loop_remove_conn(server_conn);
        znx_conn_close_and_free(server_conn);
        return;
    }

    ret = server_conn->operator->write(server_conn, data, data_len);
    if ((size_t)ret != data_len) {
        ZNX_DEF_LOG_FATAL(NULL, "send failed, expect %d, but got %d", data_len, ret);
    }

    znx_event_loop_enable_read(server_conn);
}


static void
tcp_server_accept_handler(znx_event_t *event)
{
    znx_conn_t *server_conn = event->data;

    server_conn->read.handler = znx_test_server_read;

    znx_event_loop_start_io(server_conn);
    // enabled read
    znx_event_loop_enable_read(server_conn);
}


static void client_conn_free(znx_conn_t *conn)
{
    if (conn->read.timedout) {
        ZNX_DEF_LOG_DEBUG(NULL, "client conn %p free callback, timedout", conn);
    } else {
        char buf[4];
        znx_result_t ret;
        ret = conn->operator->read(conn, buf, 4);
        if (ret == 0) {
            ZNX_DEF_LOG_DEBUG(NULL, "client conn %p free callback, server close", conn);
        } else if (ret < 0 && ret != ZNX_AGAIN) {
            ZNX_DEF_LOG_DEBUG(NULL, "client conn %p free callback, errno: %d", conn, -ret);
        }
    }

    ZNX_DEF_LOG_DEBUG(NULL, "client_conn_free: client_conn close");
    znx_event_loop_remove_conn(conn);
    znx_conn_close_and_free(conn);
}


static void test_get_and_put(znx_str_t *address_str)
{
    znx_address_t   address;
    int             port = 0;
    bzero(&address, sizeof(znx_address_t));
    znx_address_parser(address_str->data, address_str->len, &address, &port);

    znx_conn_t *conn =
        znx_t_session_pool_get(address_str, &address);
    if (conn == NULL) {
        ZNX_DEF_LOG_DEBUG(NULL, "client session get %V failed", address_str);
        return;
    } else {
        ZNX_DEF_LOG_DEBUG(NULL, "client session get %V success, conn: %p", address_str, conn);
    }

    int timeout_ms = 50;
    int64_t r = znx_random_n(2);
    if (r == 1) {
        timeout_ms = 200;
    }

    znx_str_t *peer_address_text =
        znx_conn_peer_address_text(conn);

    znx_t_session_pool_put(peer_address_text, conn, timeout_ms, client_conn_free);

    ZNX_DEF_LOG_DEBUG(NULL, "session_pool_put: %p, address: %V",
        conn, address_str);
}


static void client_timer_handler(znx_event_t *event)
{
    znx_thread_free(event);
    znx_str_t address;
    
    ZNX_STR_SET(&address, "127.0.0.1:8888");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8889");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8889");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8888");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8888");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8889");
    test_get_and_put(&address);
}


static void client_read(znx_event_t *event)
{
    znx_conn_t *conn;
    conn = event->data;
    znx_result_t ret;
    char buf[512];

    if (event->timedout) {
        znx_event_loop_remove_conn(conn);
        znx_conn_close_and_free(conn);
        return;
    }

    while (1) {
        ret = conn->operator->read(conn, buf, 512);
        if (ret > 0) {
            buf[ret] = '\0';
            ZNX_DEF_LOG_DEBUG(NULL, "test_client read data: %s", buf);
            break;
        }

        if (ret == ZNX_AGAIN) {
            znx_event_loop_enable_read(conn);
            return;
        }

        ZNX_DEF_LOG_DEBUG(NULL, "test_client read failed");
        znx_event_loop_remove_conn(conn);
        znx_conn_close_and_free(conn);
        return;
    }


    znx_str_t *peer_address_text =
        znx_conn_peer_address_text(conn);

    znx_t_session_pool_put(peer_address_text, conn, 100, client_conn_free);

    znx_event_t *e = znx_thread_calloc(sizeof(znx_event_t));
    e->handler = client_timer_handler;

    znx_event_loop_add_timer(e, 10);
    return;
}


void connect_handler(void *user_data, znx_conn_t *conn, znx_result_t res)
{
    if (res != ZNX_SUCCESS) {
        znx_conn_close_and_free(conn);
        return;
    }

    conn->read.handler = client_read;
    conn->data = user_data;

    znx_result_t ret;
    ret = conn->operator->write(conn, "aaaaa", 5);
    if (ret != 5) {
        ZNX_DEF_LOG_FATAL(NULL, "client write failed");
    }

    znx_event_loop_add_timer(&conn->read, 5);
    znx_event_loop_enable_read(conn);
}


static void udp_conn_free(znx_conn_t *conn)
{
    ZNX_DEF_LOG_DEBUG(NULL, "udp_conn close");
    znx_event_loop_remove_conn(conn);
    znx_conn_close_and_free(conn);
}


static void test_get_and_put_udp_conn()
{
    u_char          buf[128];
    znx_conn_t      *conn;
    znx_str_t       key;
    u_char          *end;
    znx_address_t   address;
    int             port = 0;

    end = znx_snprintf(buf, 512, "UDP-127.0.0.1:8888");
    key.data = buf;
    key.len = (size_t)(end - buf);
    bzero(&address, sizeof(znx_address_t));
    znx_address_parser((u_char *)"127.0.0.1:8888", strlen("127.0.0.1:8888"), &address, &port);
    conn = znx_t_session_pool_get(&key, &address);
    if (conn == NULL) {
        ZNX_DEF_LOG_DEBUG(NULL, "get udp session %v failed", &key);
    } else {
        ZNX_DEF_LOG_DEBUG(NULL, "get udp session %v success", &key);
        znx_t_session_pool_put(&key, conn, 200, udp_conn_free);
    }

    end = znx_snprintf(buf, 512, "UDP-127.0.0.1:8889");
    key.data = buf;
    key.len = (size_t)(end - buf);
    bzero(&address, sizeof(znx_address_t));
    znx_address_parser((u_char *)"127.0.0.1:8889", strlen("127.0.0.1:8889"), &address, &port);
    conn = znx_t_session_pool_get(&key, &address);
    if (conn == NULL) {
        ZNX_DEF_LOG_DEBUG(NULL, "get udp session %v failed", &key);
    } else {
        ZNX_DEF_LOG_DEBUG(NULL, "get udp session %v success", &key);
        znx_t_session_pool_put(&key, conn, 200, udp_conn_free);
    }

}


void udp_connect_handler(void *data, znx_conn_t *conn, znx_result_t res)
{
    if (res != ZNX_SUCCESS) {
        ZNX_DEF_LOG_ERROR(NULL, "udp connect failed");
        return;
    }

    znx_str_t *peer_address_text =
        znx_conn_peer_address_text(conn);

    u_char buf[128];
    u_char *end = znx_snprintf(buf, 128, "UDP-%v", peer_address_text);
    znx_str_t key;
    key.data = buf;
    key.len = (size_t)(end - buf);

    znx_t_session_pool_put(&key, conn, 100, udp_conn_free);

    test_get_and_put_udp_conn();
}


znx_bool_t tcp_client_task(void *data)
{
    znx_address_t address;

    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8888", strlen("127.0.0.1:8888"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL(NULL, "parser %s failed", "127.0.0.1:8888");
    }

    znx_conn_t *conn_1 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_1, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_1, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_2 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_2, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_2, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_3 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_3, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_3, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_4 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_4, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_4, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_5 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_5, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_5, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_6 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_6, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_6, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_7 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_7, ZNX_UDP_CONN, ZNX_CONN_EPOLL_ET);
    znx_udp_connect_setup(conn_7, &address, NULL, connect_handler, (void *)(intptr_t)1);

    znx_conn_t *conn_8 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_8, ZNX_UDP_CONN, ZNX_CONN_EPOLL_ET);
    znx_udp_connect_setup(conn_8, &address, NULL, connect_handler, (void *)(intptr_t)1);

    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8889", strlen("127.0.0.1:8889"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL(NULL, "parser %s failed", "127.0.0.1:8889");
    }

    znx_conn_t *conn_9 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_9, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_9, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)2);

    znx_conn_t *conn_10 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_10, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_10, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)2);

    znx_conn_t *conn_11 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_11, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_11, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)2);

    znx_conn_t *conn_12 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_12, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_12, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)2);

    znx_conn_t *conn_13 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_13, ZNX_TCP_CONN, ZNX_CONN_EPOLL_ET);
    znx_tcp_connect_setup(conn_13, &address, NULL, NULL, 10, connect_handler, (void *)(intptr_t)2);

    znx_conn_t *conn_14 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_14, ZNX_UDP_CONN, ZNX_CONN_EPOLL_ET);
    znx_udp_connect_setup(conn_14, &address, NULL, connect_handler, (void *)(intptr_t)2);

    znx_conn_t *conn_15 = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn_15, ZNX_UDP_CONN, ZNX_CONN_EPOLL_ET);
    znx_udp_connect_setup(conn_15, &address, NULL, connect_handler, (void *)(intptr_t)2);

    return ZNX_TRUE;
}

int main()
{
    znx_srandom((unsigned int)get_cache_time());

    // 1. 先调用proc_init
    znx_proc_init();

    // 2. core conf初始化
    znx_core_conf_create_default();
    g_core_conf->worker_threads = 1;
    g_core_conf->tcp_accepter_threads = 0;
    g_core_conf->g_mem_freelists_level = 1;
    g_core_conf->g_aligned_mem_freelists_level = 1;
    g_core_conf->worker_mem_freelist_level = 1;
    g_core_conf->worker_aligned_mem_freelists_level = 1;
    g_core_conf->thread_session_pool_size = 100;

    // 3. 初始化日志相关配置
    znx_def_log_init(0, LOG_LEVEL_DEBUG, NULL, LOG_ROLLING_DAY);

    // 4. listener conf初始化
    // 5. http conf 初始化

    znx_worker_processor_add_init_hook(tcp_client_task);

    znx_address_t address;
    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8888", strlen("127.0.0.1:8888"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL(NULL, "parser %s failed", "127.0.0.1:8888");
    }

    znx_listener_t *ls = znx_listener_alloc();
    ls->listen_addr = "127.0.0.1:8888";
    ls->address = address;
    ls->backlog = 1024;
    ls->nodelay = 1;
    ls->reuseaddr = 1;
    ls->accept_handler = tcp_server_accept_handler;

    if (!znx_listener_open(ls, INVALID_FD)) {
        ZNX_DEF_LOG_FATAL(NULL, "open listener 127.0.0.1:8888 failed");
    }

    znx_rwlock_wlock(&g_listeners->rwlock);
    znx_queue_insert_tail(&g_listeners->head, &ls->queue);
    g_listeners->version++;
    znx_rwlock_wunlock(&g_listeners->rwlock);

    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8889", strlen("127.0.0.1:8889"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL(NULL, "parser %s failed", "127.0.0.1:8889");
    }

    ls = znx_listener_alloc();
    ls->listen_addr = "127.0.0.1:8889";
    ls->address = address;
    ls->backlog = 1024;
    ls->nodelay = 1;
    ls->reuseaddr = 1;
    ls->accept_handler = tcp_server_accept_handler;

    if (!znx_listener_open(ls, INVALID_FD)) {
        ZNX_DEF_LOG_FATAL(NULL, "open listener 127.0.0.1:8888 failed");
    }

    znx_rwlock_wlock(&g_listeners->rwlock);
    znx_queue_insert_tail(&g_listeners->head, &ls->queue);
    g_listeners->version++;
    znx_rwlock_wunlock(&g_listeners->rwlock);

    znx_proc_start();

    sleep(2);

    znx_proc_shutdown();


    return 0;
}
